package com.klondike.common.concurrent;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

public class SimpleMutuxPool<T> implements MutuxPool<T> {

    private Segment[] segments;

    private long lockMaxWaitTimeout;

    private TimeUnit timeUnit;

    private int segmentShift;

    private int segmentMask;

    private long lockAutoExpriryTimeout;

    public SimpleMutuxPool(long lockAutoExpriryTimeout, long lockMaxWaitTimeout, TimeUnit timeUnit,
                           int initialCapacity, int maxCapacity, int concurrencyLevel) {
        this.lockMaxWaitTimeout = lockMaxWaitTimeout;
        this.timeUnit = timeUnit;

        this.lockAutoExpriryTimeout = timeUnit.toMillis(lockAutoExpriryTimeout);

        initSegments(initialCapacity, maxCapacity, concurrencyLevel);
    }

    public SimpleMutuxPool(long lockAutoExpriryTimeout, long lockMaxWaitTimeout, TimeUnit timeUnit) {
        this(lockAutoExpriryTimeout, lockMaxWaitTimeout, timeUnit, 16, 64, 16);
    }

    public SimpleMutuxPool() {
        this(10, 15, TimeUnit.MINUTES);
    }

    @SuppressWarnings({ "unchecked", "rawtypes" })
    private void initSegments(int initialCapacity, int maxCapacity, int concurrencyLevel) {
        // Find power-of-two sizes best matching arguments
        int segmentShift = 0;
        int segmentSize = 1;
        while (segmentSize < concurrencyLevel) {
            ++segmentShift;
            segmentSize <<= 1;
        }

        this.segmentShift = 32 - segmentShift;
        this.segmentMask = segmentSize - 1;

        int c = initialCapacity / segmentSize;
        if (c * segmentSize < initialCapacity)
            ++c;
        int cap = 1;
        while (cap < c)
            cap <<= 1;

        Segment[] segs = new SimpleMutuxPool.Segment[segmentSize];

        for (int i = 0; i < segmentSize; ++i) {
            segs[i] = new SimpleMutuxPool.Segment(cap, maxCapacity);
        }

        this.segments = segs;
    }

    private Segment segmentFor(int h) {
        return segments[(h >>> segmentShift) & segmentMask];
    }

    private int hash(T obj) {
        int h = obj.hashCode();
        // rehash based on variant of single-word Wang/Jenkins hash method
        // to reduce the hash collision.
        h += (h << 15) ^ 0xffffcd7d;
        h ^= (h >>> 10);
        h += (h << 3);
        h ^= (h >>> 6);
        h += (h << 2) + (h << 14);
        h ^= (h >>> 16);

        return h;
    }

    @Override
    public void lockFor(T obj) {
        int h = hash(obj);

        boolean locked = segmentFor(h).tryLockFor(h, obj, lockMaxWaitTimeout, timeUnit);
        if (!locked) {
            String msg = String.format("Could not get locked for %s in %d %s", obj, lockMaxWaitTimeout, timeUnit);
            throw new IllegalStateException(msg);
        }

    }

    @Override
    public boolean tryLockFor(T obj, long timeout, TimeUnit unit) {
        int h = hash(obj);

        return segmentFor(h).tryLockFor(h, obj, timeout, unit);
    }

    @Override
    public boolean isLockedFor(T obj) {
        int h = hash(obj);
        return segmentFor(h).isLockedFor(h, obj);
    }

    @Override
    public void unlockFor(T obj) {
        int h = hash(obj);
        segmentFor(h).unlockedFor(h, obj);
    }

    private static class Bucket<T> {
        int h;
        T key;
        volatile Bucket<T> next;
        volatile Mutux value;

        Bucket<T> before, after;

        public Bucket(int h, T key, Mutux value, Bucket<T> next) {
            this.h = h;
            this.key = key;
            this.value = value;
            this.next = next;
        }

        @Override
        public String toString() {
            return "h=" + h + ", " + key + "=" + value;
        }

        public void addBefore(Bucket<T> given) {
            this.after = given;
            this.before = given.before;
            before.after = this;
            after.before = this;
        }

        public void remove() {
            before.after = after;
            after.before = before;
        }

        public void recordAccess(Bucket<T> given) {
            remove();
            addBefore(given);
        }
    }

    private class Segment extends ReentrantLock {

        private static final long serialVersionUID = 1L;

        private int threshold;

        private Bucket<T>[] buckets;

        private Bucket<T> head;

        private int count;

        @SuppressWarnings("unchecked")
        public Segment(int capacity, int threshold) {
            buckets = new Bucket[capacity];
            this.threshold = threshold;
            head = new Bucket<T>(0, null, null, null);
            head.before = head.after = head;
        }

        private Mutux getOrAllocateMutuxFor(int h, T key, boolean allocate) {
            lock();
            try {
                Mutux mutux = null;

                int index = bucketIndexFor(h);
                Bucket<T> bucketHead = buckets[index];
                Bucket<T> e = findBucket(h, key, bucketHead);

                if (allocate) {
                    if (e == null) {
                        removeEldestEntryIfCan();

                        mutux = allocateMutux();

                        buckets[index] = e = new Bucket<T>(h, key, mutux, buckets[index]);

                        count++;
                        e.addBefore(head);

                        return mutux;
                    } else if (discardMutux(e.value)) {
                        e.value = allocateMutux();
                    }
                }

                if (e != null) {
                    e.recordAccess(head);
                    mutux = e.value;
                }

                return mutux;
            } finally {
                unlock();
            }
        }

        private int bucketIndexFor(int h) {
            return h & (buckets.length - 1);
        }

        private Bucket<T> findBucket(int h, T key, Bucket<T> e) {
            while (e != null && (e.h != h || !key.equals(e.key))) {
                e = e.next;
            }
            return e;
        }

        private Bucket<T> removeEldestEntryIfCan() {
            Bucket<T> e = head;
            Bucket<T> removed = null;
            if (count < threshold) {
                if ((e = e.after) != head) {
                    Mutux mutux = e.value;
                    if (discardMutux(mutux)) {
                        removeBucket(e);
                        removed = e;
                    }
                }
            } else {
                while ((e = e.after) != head) {
                    Mutux mutux = e.value;
                    if (discardMutux(mutux)) {
                        removeBucket(e);
                        removed = e;
                    }
                }
            }

            return removed;
        }

        private void removeBucket(Bucket<T> given) {
            Bucket<T>[] b = buckets;
            int h = given.h;
            int index = h & (b.length - 1);

            Bucket<T> last = b[index];
            Bucket<T> e = last;
            while (e != null) {
                if (e == given) {
                    count--;

                    if (b[index] == e) {
                        b[index] = e.next;
                    } else {
                        last.next = e.next;
                    }

                    e.remove();
                    break;
                }
                last = e;
                e = e.next;
            }
        }

        private boolean discardMutux(Mutux mutux) {
            return isLockAutoExpiredTimeout(mutux);
        }

        private boolean isLockAutoExpiredTimeout(Mutux mutux) {
            return System.currentTimeMillis() - mutux.wasLockedOn() >= lockAutoExpriryTimeout;
        }

        public void unlockedFor(int h, T obj) {
            Mutux mutux = getOrAllocateMutuxFor(h, obj, false);

            if (mutux == null) {
                throw new IllegalMonitorStateException();
            }

            mutux.unlock();
        }

        public boolean isLockedFor(int h, T obj) {
            Mutux mutux = getOrAllocateMutuxFor(h, obj, false);

            return mutux == null || discardMutux(mutux) ? false : mutux.isLocked();
        }

        public boolean tryLockFor(int h, T obj, long timeout, TimeUnit timeUnit) {
            try {
                return getOrAllocateMutuxFor(h, obj, true).tryLock(timeout, timeUnit);
            } catch (InterruptedException e) {
                return false;
            }
        }

        public void clear() {
            lock();
            try {
                Segment segment = this;
                
                Bucket<T> head = segment.head;
                head.after = head.before = head;
                Bucket<T>[] buckets = segment.buckets;

                for (int i = 0; i < buckets.length; i++) {
                    buckets[i] = null;
                }
                
            } finally {
                unlock();
            }
        }
    }

    protected Mutux allocateMutux() {
        return new MutuxImpl();
    }

    public void clear() {
        Segment[] segs = segments;

        for (Segment segment : segs) {
            segment.clear();
        }
    }

}
